package org.teresoft.collections.concurrent;

import static java.lang.Integer.bitCount;
import static java.lang.System.arraycopy;
import static org.teresoft.math.Arithmetic.ceilingPowerOf2;

import java.util.AbstractCollection;
import java.util.AbstractSet;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicReferenceArray;

import org.teresoft.Preconditions;
import org.teresoft.collections.Out;

@SuppressWarnings({"unchecked", "rawtypes"})
public class ConcurrentHashTrie<K, V> implements ConcurrentMap<K, V> {

    private final static Out<Object> NULL_OUT = new Out<Object>();

    private final AtomicReferenceArray<Node> rootTable;
    private final int rootTableMask;

    public ConcurrentHashTrie(int size) {
        int rootTableSize = ceilingPowerOf2(size);
        rootTable = new AtomicReferenceArray<Node>(rootTableSize);
        rootTableMask = rootTableSize - 1;
    }
    
    @Override
    public V putIfAbsent(K key, V value) {
        Preconditions.checkNotNull(key, "key must not be null");
        Preconditions.checkNotNull(value, "value must not be null");
        
        final int hashCode = hash(key.hashCode());
        final int rootHash = rootHash(hashCode);
        
        Node existing;
        Node update;
        do {
            existing = rootTable.get(rootHash); // Snapshot the current node.
            if (existing != null) {
                Object result = existing.get(hashCode, key);
                if (result != null) {
                    return (V) result;
                }
                update = existing.put(0, hashCode, key, value, NULL_OUT);
            } else {
                update = new LeafNode(hashCode, key, value);
            }
        } while (!rootTable.compareAndSet(rootHash, existing, update));
        
        return null;
    }

    @Override
    public boolean remove(Object key, Object value) {
        Preconditions.checkNotNull(key, "key must not be null");
        Preconditions.checkNotNull(value, "value must not be null");
        
        final int hashCode = hash(key.hashCode());
        final int rootHash = rootHash(hashCode);
        
        Node existing;
        Node update;
        do {
            existing = rootTable.get(rootHash); // Snapshot the current node.
            if (existing != null) {
                Object oldValue = existing.get(hashCode, key);
                if (value == oldValue || value.equals(oldValue)) {
                    update = existing.remove(hashCode, key, NULL_OUT);
                } else {
                    return false;
                }
            } else {
                return false;
            }
        } while (!rootTable.compareAndSet(rootHash, existing, update));
        
        return true;
    }

    @Override
    public V replace(K key, V value) {
        Preconditions.checkNotNull(key, "key must not be null");
        Preconditions.checkNotNull(value, "value must not be null");
        
        final int hashCode = hash(key.hashCode());
        final int rootHash = rootHash(hashCode);
        
        Node existing;
        Node update;
        Object oldValue = null;
        do {
            existing = rootTable.get(rootHash); // Snapshot the current node.
            if (existing != null) {
                oldValue = existing.get(hashCode, key);
                if (null != oldValue) {
                    update = existing.put(0, hashCode, key, value, NULL_OUT);
                } else {
                    return null;
                }
            } else {
                return null;
            }
        } while (!rootTable.compareAndSet(rootHash, existing, update));
        
        return (V) oldValue;
    }

    @Override
    public boolean replace(K key, V oldValue, V newValue) {
        Preconditions.checkNotNull(key, "key must not be null");
        Preconditions.checkNotNull(oldValue, "oldValue must not be null");
        Preconditions.checkNotNull(newValue, "newValue must not be null");
        
        final int hashCode = hash(key.hashCode());
        final int rootHash = rootHash(hashCode);
        
        Node existing;
        Node update;
        do {
            existing = rootTable.get(rootHash); // Snapshot the current node.
            if (existing != null) {
                Object existingValue = existing.get(hashCode, key);
                if (existingValue == oldValue || existingValue.equals(oldValue)) {
                    update = existing.put(0, hashCode, key, newValue, NULL_OUT);
                } else {
                    return false;
                }
            } else {
                return false;
            }
        } while (!rootTable.compareAndSet(rootHash, existing, update));
        
        return true;
    }

    @Override
    public void clear() {
        for (int i = 0, n = rootTable.length(); i < n; i++) {
            rootTable.set(i, null);
        }
    }

    @Override
    public boolean containsKey(Object key) {
        return get(key) != null;
    }

    @Override
    public boolean containsValue(Object value) {
        Preconditions.checkNotNull(value, "value must not be null");
        
        return values().contains(value);
    }

    @Override
    public V get(Object key) {
        Preconditions.checkNotNull(key, "key must not be null");
        
        final int hashCode = hash(key.hashCode());
        Node node = rootTable.get(rootHash(hashCode));
        return (node != null) ? (V) node.get(hashCode, key) : null;
    }

    @Override
    public boolean isEmpty() {
        return 0 == size();
    }

    @Override
    public V put(K key, V value) {
        Preconditions.checkNotNull(key, "key must not be null");
        Preconditions.checkNotNull(value, "key must not be null");
        
        final int hashCode = hash(key.hashCode());
        final int rootHash = rootHash(hashCode);
        
        Node node = rootTable.get(rootHash);
        V result = null;
        
        if (node != null) {
            Out<Object> out = new Out<Object>();
            rootTable.set(rootHash, node.put(0, hashCode, key, value, out));
            result = (V) out.value;
        } else {
            rootTable.set(rootHash, new LeafNode(hashCode, key, value));
        }
        
        return result;
    }

    @Override
    public void putAll(Map<? extends K, ? extends V> m) {
        for (Map.Entry<? extends K, ? extends V> entry : m.entrySet()) {
            put(entry.getKey(), entry.getValue());
        }
    }

    @Override
    public V remove(Object key) {
        final int hashCode = hash(key.hashCode());
        final int rootHash = rootHash(hashCode);
        
        Node node = rootTable.get(rootHash);
        if (node != null) {
            Out<Object> out = new Out<Object>();
            rootTable.set(rootHash, node.remove(hashCode, key, out));
            return (V) out.value;
        }
        
        return null;
    }

    @Override
    public int size() {
        int size = 0;
        for (int i = 0, n = rootTable.length(); i < n; i++) {
            Node node = rootTable.get(i);
            if (node != null) {
                size += node.size();
            }
        }
        return size;
    }
    
    private final int rootHash(final int hashCode) {
        return hashCode & rootTableMask;
    }

    /**
     * Stolen from java.util.HashMap.
     * 
     * @param h
     * @return
     */
    static int hash(int h) {
        // This function ensures that hashCodes that differ only by
        // constant multiples at each bit position have a bounded
        // number of collisions (approximately 8 at default load factor).
        h ^= (h >>> 20) ^ (h >>> 12);
        return h ^ (h >>> 7) ^ (h >>> 4);
    }
    
    // ------------------ Trie Nodes ----------------------------
    
    interface Node extends Iterable {
        int getHashCode();
        Object get(final int hashCode, Object key);
        Node put(final int level, final int hashCode, Object key, Object value, Out<Object> out);
        Node remove(final int hashCode, Object key, Out<Object> out);
        int size();
    }
    
    static final class BitMappedNode implements Node {
        
        private final int level;
        private final Node[] table;
        private final int bitMap;
        private final int size;
        
        private BitMappedNode(int level, Node[] table, int bitMap, int size) {
            this.level = level;
            this.table = table;
            this.bitMap = bitMap;
            this.size = size;
        }
        
        @Override
        public Object get(final int hashCode, Object key) {
            int bitIndex = 1 << getPartialHash(hashCode, level);
            if ((bitIndex & bitMap) != 0) {
                int index = bitCount((bitIndex - 1) & bitMap);
                return table[index].get(hashCode, key);
            }
            return null;
        }
        
        @Override
        public BitMappedNode put(int level, final int hashCode, Object key, Object value, Out<Object> out) {
            int bitIndex = 1 << getPartialHash(hashCode, level);
            int index = bitCount((bitIndex -1) & bitMap);
            
            if ((bitIndex & bitMap) == 0) {
                LeafNode node = new LeafNode(hashCode, key, value);
                return new BitMappedNode(level, copyAndInsert(table, node, index), bitMap | bitIndex, this.size + 1); 
            } else {
                Node node = table[index].put(level + 1, hashCode, key, value, out);
                int size = (out.value == null) ? this.size + 1 : this.size;
                return new BitMappedNode(level, copyAndReplace(table, node, index), bitMap, size); 
            }
        }
        
        @Override
        public Node remove(final int hashCode, Object key, Out<Object> out) {
            int bitIndex = 1 << getPartialHash(hashCode, level);
            
            if ((bitIndex & bitMap) != 0) {
                int index = bitCount((bitIndex - 1) & bitMap);
                Node child = table[index].remove(hashCode, key, out);
                if (out.value != null) {
                    if (child == null) {
                        if (table.length != 2) {                    
                            return new BitMappedNode(level, copyAndRemove(table, index), bitMap ^ bitIndex, size - 1);
                        } else {
                            return table[index ^ 1];
                        }
                    } else {
                        return new BitMappedNode(level, copyAndReplace(table, child, index), bitMap, size - 1);
                    }
                }
            }
            
            return this;
        }
        
        @Override
        public int size() {
            return size;
        }
        
        private static int getPartialHash(final int hashCode, int level) {
            int levelShift = level * 5;
            int mask = 0x1F << levelShift;
            return ((mask & hashCode) >>> levelShift);
        }

        public static BitMappedNode create(final int hashCode1, Object key1, Object value1,
                                             final int hashCode2, Object key2, Object value2) {
            return BitMappedNode.create(0, new LeafNode(hashCode1, key1, value1), new LeafNode(hashCode2, key2, value2));
        }
        
        @Override
        public int getHashCode() {
            return table[0].getHashCode();
        }

        private static BitMappedNode create(int level, Node nodeA, Node nodeB) {
            int partialHashA = getPartialHash(nodeA.getHashCode(), level);
            int partialHashB = getPartialHash(nodeB.getHashCode(), level);
            
            Node[] table;
            if (partialHashA < partialHashB) {
                table = new Node[] { nodeA, nodeB };
            } else if (partialHashA > partialHashB) {
                table = new Node[] { nodeB, nodeA };
            } else if (nodeA.hashCode() != nodeB.hashCode()) {
                Node node = create(level + 1, nodeA, nodeB);
                table = new Node[] { node };
            } else {
                throw new UnsupportedOperationException();
            }
            return new BitMappedNode(level, table, (1 << partialHashA) | (1 << partialHashB), nodeA.size() + nodeB.size());
        }
        
        @Override
        public Iterator iterator() {
            return new TableIterator(table);
        }
    }
        
    static final class LeafNode implements Node, Iterable, Map.Entry {
        
        private final int hashCode;
        private final Object key;
        private final Object value;

        public LeafNode(final int hashCode, Object key, Object value) {
            this.hashCode = hashCode;
            this.key = key;
            this.value = value;
        }

        @Override
        public Object get(final int hashCode, Object key) {
            if (keyMatches(hashCode, key)) {
                return value;
            }
            return null;
        }

        @Override
        public Node put(int level, final int hashCode, Object key, Object value, Out<Object> out) {
            Node node = new LeafNode(hashCode, key, value);
            if (this.hashCode == hashCode) {
                if (this.key == key || this.key.equals(key)) {
                    out.value = this.value;
                    return node;
                } else {
                    return ListNode.create(this, node);
                }
            } else {
                return BitMappedNode.create(level, this, node);
            }
        }
        
        @Override
        public Node remove(final int hashCode, Object key, Out<Object> out) {
            if (keyMatches(hashCode, key)) {
                out.value = this.value;
                return null;
            } else {
                return this;
            }
        }
        
        private boolean keyMatches(final int hashCode, Object key) {
            return this.hashCode == hashCode && (this.key == key || this.key.equals(key));
        }
        
        @Override
        public Object getKey() {
            return key;
        }
        
        @Override
        public Object getValue() {
            return value;
        }
        
        @Override
        public int getHashCode() {
            return hashCode;
        }        
        
        @Override
        public int size() {
            return 1;
        }
        
        @Override
        public Iterator<Map.Entry> iterator() {
            return new Iterator<Map.Entry>() {
                
                boolean hasNext = true;
                
                @Override
                public Map.Entry next() {
                    if (hasNext) {
                        hasNext = false;
                        return LeafNode.this;
                    } else {
                        throw new NoSuchElementException();
                    }
                }
                
                @Override
                public boolean hasNext() {
                    return hasNext;
                }
                
                @Override
                public void remove() {
                    throw new UnsupportedOperationException();
                }
            };
        }
        
        @Override
        public Object setValue(Object value) {
            throw new UnsupportedOperationException();
        }
    }
    
    static final class ListNode implements Node {
        
        private final Node[] table;

        public ListNode(Node[] table) {
            this.table = table;
        }
        
        @Override
        public Object get(final int hashCode, Object key) {
            for (int i = 0, n = table.length; i < n; i++) {
                Object value = table[i].get(hashCode, key);
                if (value != null) {
                    return value;
                }
            }
            return null;
        }

        @Override
        public ListNode put(int level, final int hashCode, Object key, Object value, Out<Object> out) {
            int index = find(hashCode, key, out);
            
            if (out.value == null) {
                return new ListNode(copyAndInsert(table, new LeafNode(hashCode, key, value), table.length));
            } else {
                return new ListNode(copyAndReplace(table, new LeafNode(hashCode, key, value), index));
            }
        }

        @Override
        public Node remove(final int hashCode, Object key, Out<Object> out) {
            int index = find(hashCode, key, out);
            
            if (out.value != null) {
                if (table.length != 2) {
                    return new ListNode(copyAndRemove(table, index));
                } else {                    
                    return table[index ^ 1];
                }
            }
            
            return this;
        }
        
        @Override
        public int getHashCode() {
            return table[0].getHashCode();
        }
        
        @Override
        public int size() {
            return table.length;
        }
        
        @Override
        public Iterator iterator() {
            return new TableIterator(table);
        }

        private int find(final int hashCode, Object key, Out<Object> out) {
            for (int i = 0, n = table.length; i < n; i++) {
                Object existingValue = table[i].get(hashCode, key);
                if (existingValue != null) {
                    out.value = existingValue;
                    return i;
                }
            }
            return 0;
        }
        
        public static ListNode create(Node nodeA, Node nodeB) {
            return new ListNode(new Node[] { nodeA, nodeB });
        }
    }
    
    // -------------- Iterators -----------------------
    
    @Override
    public Set<java.util.Map.Entry<K, V>> entrySet() {
        return new AbstractSet<Entry<K,V>>() {

            @Override
            public Iterator<java.util.Map.Entry<K, V>> iterator() {
                return new RootTableIterator();
            }

            @Override
            public int size() {
                return ConcurrentHashTrie.this.size();
            }
        };
    }    
    
    @Override
    public Collection<V> values() {
        return new AbstractCollection<V>() {

            @Override
            public Iterator<V> iterator() {
                return new Iterator<V>() {

                    private final RootTableIterator rootTableIterator = new RootTableIterator();
                    
                    @Override
                    public boolean hasNext() {
                        return rootTableIterator.hasNext();
                    }

                    @Override
                    public V next() {
                        return rootTableIterator.next().getValue();
                    }

                    @Override
                    public void remove() {
                        rootTableIterator.remove();
                    }
                };
            }

            @Override
            public int size() {
                return ConcurrentHashTrie.this.size();
            }
        };
    }

    @Override
    public Set<K> keySet() {
        return new AbstractSet<K>() {

            @Override
            public Iterator<K> iterator() {
                return new Iterator<K>() {

                    private final RootTableIterator rootTableIterator = new RootTableIterator();

                    @Override
                    public boolean hasNext() {
                        return rootTableIterator.hasNext();
                    }

                    @Override
                    public K next() {
                        return rootTableIterator.next().getKey();
                    }

                    @Override
                    public void remove() {
                        rootTableIterator.remove();
                    }
                    
                };
            }

            @Override
            public int size() {
                return ConcurrentHashTrie.this.size();
            }
        };
    }

    private final class RootTableIterator implements Iterator<Entry<K, V>> {
        private int pointer = 0;
        private Iterator<Map.Entry> child = null;

        @Override
        public boolean hasNext() {
            
            while ((child == null || !child.hasNext()) && pointer < rootTable.length()) {
                Node node = rootTable.get(pointer++);
                if (node != null) {
                    child = node.iterator();
                }
            }
            
            return child != null ? child.hasNext() : false;
        }

        @Override
        public java.util.Map.Entry<K, V> next() {
            
            if (!hasNext()) {
                throw new NoSuchElementException();
            }
            
            return child.next();
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }
    }    
    
    private static final class TableIterator implements Iterator<Map.Entry> {
        private final Node[] table;
        private int pointer = -1;
        private Iterator<Map.Entry> child = null;
        
        public TableIterator(Node[] table) {
            this.table = table;
        }

        @Override
        public boolean hasNext() {

            if (pointer >= table.length) {
                return false;
            }

            if (child == null) {
                child = table[++pointer].iterator();
            }

            return child.hasNext();
        }

        @Override
        public Map.Entry next() {

            if (pointer >= table.length) {
                throw new NoSuchElementException();
            }

            if (child == null) {
                child = table[++pointer].iterator();
            }

            Map.Entry result = child.next();

            if (!child.hasNext()) {
                child = null;
            }

            return result;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }
    }    
    
    // -------------- Utility Methods ----------------------

    public static Node[] copyAndRemove(Node[] nodes, int index) {
        int len = nodes.length - 1;
        Node[] newNodes = new Node[len];
        arraycopy(nodes, 0, newNodes, 0, index);
        arraycopy(nodes, index + 1, newNodes, index, len - index);
        return newNodes;
    }

    public static Node[] copyAndReplace(Node[] nodes, Node node, int index) {
        int len = nodes.length;
        
        Node[] newNodes = new Node[len];
        System.arraycopy(nodes, 0, newNodes, 0, len);
        newNodes[index] = node;
        return newNodes;
    }

    public static Node[] copyAndInsert(Node[] nodes, Node node, int index) {
        int len = nodes.length;
    
        Node[] newNodes = (Node[]) new Node[len + 1];
        arraycopy(nodes, 0, newNodes, 0, index);
        arraycopy(nodes, index, newNodes, index + 1, len - index);
        newNodes[index] = node;
        return newNodes;
    }
}
